#include <iostream>
#include <fstream>

#include "TROOT.h"
#include "TFile.h"
#include "TTree.h"
#include "TSystem.h"
#include "TDataType.h"

#include "TLorentzVector.h"
#include "HepMC/IO_GenEvent.h"
#include "HepMC/GenEvent.h"
#include "HepMC/GenParticle.h"
#include "HepMC/GenVertex.h"


using namespace std;
using namespace HepMC;

int main(int narg, char** argv)
{

  cout<<"Usage: ./converter input output"<<endl;
  
  TFile out(argv[2], "RECREATE");
  cout<<"start readin"<<endl;
  IO_GenEvent readin(argv[1], std::ios::in);

  int nskip=-1;

  GenEvent* evt;
  if(readin.rdstate() == ios::failbit)
    {cout<<"file not found"<<endl;}
  /*
  if(narg > 3)
    {
      nskip=10000*atoi(argv[3]);
      int nnow=0;
      while(nnow<nskip)
	{
	  evt=readin.read_next_event();
	  nnow++;
	}
      if(evt)
	delete evt;
    }

  int N=0, Nmax=-1;
  
  if(narg > 4)
    {
      Nmax=10000*atoi(argv[4]);
    }
  */

  evt=readin.read_next_event();
  if(evt==NULL)
    {
      cout<<"ERROR: No more events to be read...exit"<<endl;
      return 1;
    }


  TTree tree("STDHEP", "STDHEP");

  const int kMaxGenParticle=10000;

  Double_t        weight;
  Int_t           PID[kMaxGenParticle];
  Int_t           Status[kMaxGenParticle];
  Double_t        E[kMaxGenParticle];
  Double_t        Px[kMaxGenParticle];
  Double_t        Py[kMaxGenParticle];
  Double_t        Pz[kMaxGenParticle];
  Double_t        Pt[kMaxGenParticle];
  Double_t        Eta[kMaxGenParticle];
  Double_t        Phi[kMaxGenParticle];
  Double_t        Rapidity[kMaxGenParticle];
  Double_t        T[kMaxGenParticle];
  Double_t        X[kMaxGenParticle];
  Double_t        Y[kMaxGenParticle];
  Double_t        Z[kMaxGenParticle];
  Int_t           size;

  tree.Branch("weight", &weight, "weight/D");
  tree.Branch("size", &size, "size/I");


  tree.Branch("PID", PID, "PID[size]/I");
  tree.Branch("Status", Status, "Status[size]/I");
  /*
  tree.Branch("M1", M1, "M1[size]/I");
  tree.Branch("M2", M2, "M2[size]/I");
  tree.Branch("D1", D1, "D1[size]/I");
  tree.Branch("D2", D2, "D2[size]/I");
  */

  tree.Branch("E", E, "E[size]/D");
  tree.Branch("Px", Px, "Px[size]/D");
  tree.Branch("Py", Py, "Py[size]/D");
  tree.Branch("Pz", Pz, "Pz[size]/D");
  tree.Branch("Pt", Pt, "Pt[size]/D");
  tree.Branch("Eta", Eta, "Eta[size]/D");
  tree.Branch("Phi", Phi, "Phi[size]/D");

  /*
  tree.Branch("Rapidity", Rapidity, "Rapidity[size]/F");
  tree.Branch("T", T, "T[size]/F");
  tree.Branch("X", X, "X[size]/F");
  tree.Branch("Y", Y, "Y[size]/F");
  tree.Branch("Z", Z, "Z[size]/F");
  */

  int N=0;

  while(evt)
    {
      N++;
      //iterate over particles;
      GenEvent::particle_const_iterator p;
      
      size=0;
      for(p=evt->particles_begin(); 
	  p!=evt->particles_end(); ++p)
	if( (*p)->status() ==1 || abs((*p)->pdg_id())==6 )
	  {
	    if(size==kMaxGenParticle)
	      break;
	    //grab 4 vector
	    const FourVector& v=(*p)->momentum();
	    E[size]=v.e();
	    Px[size]=v.px();
	    Py[size]=v.py();
	    Pz[size]=v.pz();
	    Eta[size]=v.eta();
	    Phi[size]=v.phi();
	    Pt[size]=v.perp();
	    Status[size]=(*p)->status();
	    PID[size]=(*p)->pdg_id();
	    
	    size++;
	  }

      
      weight=evt->weights()[0];
      tree.Fill();
      
      if(N%1000==0)
	cout<<"processing event "<<N<<endl;

      delete evt;
      //read next event
      evt= readin.read_next_event();
    }

  cout<<N<<" events processed"<<endl;

  tree.Write();
  out.Close();
  return 1;
}
